package com.googlecode.recycled.drools;

import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.googlecode.recycled.lang.assertions.Assert;
import org.drools.KnowledgeBase;
import org.drools.KnowledgeBaseFactory;
import org.drools.builder.KnowledgeBuilder;
import org.drools.builder.KnowledgeBuilderFactory;
import org.drools.builder.ResourceType;
import org.drools.command.Command;
import org.drools.command.CommandFactory;
import org.drools.event.rule.WorkingMemoryEventListener;
import org.drools.io.ResourceFactory;
import org.drools.runtime.ExecutionResults;
import org.drools.runtime.StatefulKnowledgeSession;
import org.drools.runtime.StatelessKnowledgeSession;

import java.io.InputStream;
import java.util.List;
import java.util.Map;

/**
 * <p>Template Method making dealing with Drools sessions easier. It holds Drools KnowledgeBase
 * and allows to perform both stateful and stateless operations against it.</p>
 * <p/>
 * <p>DroolsTemplate is thread-safe.</p>
 */
public class DroolsTemplate {

    protected final KnowledgeBase knowledgeBase;

    protected Map<String, Object> globalServices = Maps.newHashMap();

    public DroolsTemplate(Iterable<InputStream> rules) {
        Assert.notEmpty(rules);

        KnowledgeBuilder builder = KnowledgeBuilderFactory.newKnowledgeBuilder();
        for (InputStream ruleInputStream : rules) {
            builder.add(ResourceFactory.newInputStreamResource(ruleInputStream), ResourceType.DRL);
        }
        if (builder.hasErrors()) {
            throw new RuntimeException(builder.getErrors().toString());
        }
        knowledgeBase = buildKnowledgeBase(builder);
    }

    public DroolsTemplate(Iterable<InputStream> rules, Iterable<InputStream> dsls) {
        Assert.notEmpty(rules);
        Assert.notEmpty(dsls);

        KnowledgeBuilder builder = KnowledgeBuilderFactory.newKnowledgeBuilder();
        for (InputStream dslInputStream : dsls) {
            builder.add(ResourceFactory.newInputStreamResource(dslInputStream), ResourceType.DSL);
        }
        for (InputStream ruleInputStream : rules) {
            builder.add(ResourceFactory.newInputStreamResource(ruleInputStream), ResourceType.DSLR);
        }
        knowledgeBase = buildKnowledgeBase(builder);
    }

    protected KnowledgeBase buildKnowledgeBase(KnowledgeBuilder builder) {
        if (builder.hasErrors()) {
            throw new RuntimeException(builder.getErrors().toString());
        }
        KnowledgeBase knowledgeBase = KnowledgeBaseFactory.newKnowledgeBase();
        knowledgeBase.addKnowledgePackages(builder.getKnowledgePackages());
        return knowledgeBase;
    }

    public int stateful(Object... facts) {
        Assert.notEmpty(facts);

        StatefulKnowledgeSession session = knowledgeBase.newStatefulKnowledgeSession();
        try {
            for (Map.Entry<String, Object> globalEntry : globalServices.entrySet()) {
                session.setGlobal(globalEntry.getKey(), globalEntry.getValue());
            }
            for (Object fact : facts) {
                session.insert(fact);
            }
            return session.fireAllRules();
        } finally {
            session.dispose();
        }
    }

    public StatelessSessionContext stateless() {
        return new StatelessSessionContext();
    }

    public <O> Iterable<O> transform(Object... facts) {
        Assert.notEmpty(facts);

        StatelessSessionContext statelessSessionContext = new StatelessSessionContext();

        FactHandlerListener factHandlerListener = new FactHandlerListener();
        statelessSessionContext.workingMemoryEventListeners(factHandlerListener);

        statelessSessionContext.forFacts(facts);

        return (Iterable<O>) factHandlerListener.getFactsHistory();
    }

    public <O> O transformToSingleton(Object... facts) {
        return (O) Assert.single(transform(facts));
    }

    public void setGlobalServices(Map<String, Object> globalServices) {
        this.globalServices = Assert.notNull(globalServices);
    }

    public void addGlobalService(String serviceId, Object service) {
        globalServices.put(Assert.notNull(serviceId), Assert.notNull(service));
    }

    public class StatelessSessionContext {

        protected List<WorkingMemoryEventListener> workingMemoryEventListeners = Lists.newLinkedList();

        protected List<Command> beforeInsertCommands = Lists.newLinkedList();

        protected List<Command> beforeFireCommands = Lists.newLinkedList();

        protected List<Command> afterFireCommands = Lists.newLinkedList();

        public ExecutionResults forFacts(Object... facts) {
            Assert.notEmpty(facts);

            StatelessKnowledgeSession session = knowledgeBase.newStatelessKnowledgeSession();
            for (WorkingMemoryEventListener workingMemoryEventListener : workingMemoryEventListeners) {
                session.addEventListener(workingMemoryEventListener);
            }
            for (Map.Entry<String, Object> globalEntry : globalServices.entrySet()) {
                session.setGlobal(globalEntry.getKey(), globalEntry.getValue());
            }
            List<Command> commands = Lists.newLinkedList();
            commands.addAll(beforeInsertCommands);
            commands.add(CommandFactory.newInsertElements(Lists.newArrayList(facts)));
            commands.addAll(beforeFireCommands);
            commands.add(CommandFactory.newFireAllRules());
            commands.addAll(afterFireCommands);
            return session.execute(CommandFactory.newBatchExecution(commands));
        }

        public void workingMemoryEventListeners(WorkingMemoryEventListener... workingMemoryEventListeners) {
            this.workingMemoryEventListeners.addAll(Lists.newArrayList(workingMemoryEventListeners));
        }

        public void beforeInsertCommands(Command... commands) {
            beforeInsertCommands.addAll(Lists.newArrayList(commands));
        }

        public void beforeFireCommands(Command... commands) {
            beforeFireCommands.addAll(Lists.newArrayList(commands));
        }

        public void afterFireCommands(Command... commands) {
            afterFireCommands.addAll(Lists.newArrayList(commands));
        }

    }

}
